package accountsMerge_721;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;

public class Solution {
    public List<List<String>> accountsMerge(List<List<String>> accounts) {
        HashMap<String, String> namePt = new HashMap<>();
        HashMap<String, Integer> indexPt = new HashMap<>();
        int emailCount = 0;
        for (List<String> account : accounts) {
            String name = account.get(0);
            for (int j = 1; j < account.size(); j++) {
                String email = account.get(j);
                if (!indexPt.containsKey(email)) {
                    namePt.put(email, name);
                    indexPt.put(email, emailCount++);
                }
            }
        }

        Union union = new Union(emailCount);
        for (List<String> account : accounts) {
            String firstEmail = account.get(1);
            int firstIndex = indexPt.get(firstEmail);
            int size = account.size();
            for (int i = 2; i < size; i++) {
                String nextEmail = account.get(i);
                int nextIndex = indexPt.get(nextEmail);
                union.merge(firstIndex, nextIndex);
            }
        }

        HashMap<Integer, List<String>> indexToEmails = new HashMap<Integer, List<String>>();
        for (String email : indexPt.keySet()) {
            int index = union.find(indexPt.get(email));
            List<String> account = indexToEmails.getOrDefault(index, new ArrayList<>());
            account.add(email);
            indexToEmails.put(index, account);
        }

        List<List<String>> merged = new ArrayList<List<String>>();
        for (List<String> emails : indexToEmails.values()) {
            Collections.sort(emails);
            String name = namePt.get(emails.get(0));
            List<String> account = new ArrayList<>();
            account.add(name);
            account.addAll(emails);
            merged.add(account);
        }

        return merged;
    }
}


class Union {
    private final int[] parent;

    public Union(int size) {
        parent = new int[size];
        for (int i = 0; i < parent.length; i++) {
            parent[i] = i;
        }
    }


    int find(int index) {
        while (parent[index] != index) {
            index = parent[index];
        }
        return parent[index];
    }


    void merge(int i1, int i2) {
        parent[find(i2)] = find(i1);
    }

}